package org.hepeng.workx.spring.cloud.netflix.ribbon.loadbalancer;

import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;
import lombok.Data;
import org.apache.commons.collections.CollectionUtils;
import org.hepeng.workx.spring.cloud.netflix.ribbon.RibbonRequestContext;

import java.util.List;
import java.util.Objects;

/**
 * @author he peng
 */

@Data
public class CompositeRule extends AbstractThreadIsolationRule {

    private static final ThreadLocal<AbstractThreadIsolationRule> RULE_THREAD_LOCAL = new ThreadLocal<>();
    private List<AbstractThreadIsolationRule> rules;

    public CompositeRule() {
    }

    public CompositeRule(List<AbstractThreadIsolationRule> rules) {
        this.rules = rules;
    }

    @Override
    public List<Server> getParticipateLoadBalancingServers(Object key, RibbonRequestContext ribbonContext, List<Server> originalAllServers) {
        try {
            AbstractThreadIsolationRule rule = RULE_THREAD_LOCAL.get();
            return Objects.nonNull(rule) ? rule.getParticipateLoadBalancingServers(key, ribbonContext, originalAllServers) : null;
        } finally {
            RULE_THREAD_LOCAL.remove();
        }
    }

    @Override
    protected boolean isSupported(Object key, RibbonRequestContext ribbonContext) {
        boolean supported = false;
        if (CollectionUtils.isNotEmpty(rules)) {
            for (AbstractThreadIsolationRule rule : rules) {
                supported = rule.isSupported(key, ribbonContext);
                if (supported) {
                    RULE_THREAD_LOCAL.set(rule);
                    break;
                }
            }
        }
        return supported;
    }

    @Override
    public void setLoadBalancer(ILoadBalancer lb) {
        super.setLoadBalancer(lb);
        if (CollectionUtils.isNotEmpty(rules)) {
            for (AbstractThreadIsolationRule rule : rules) {
                rule.setLoadBalancer(lb);
            }
        }
    }
}
